using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using RTMP.SimpleClient.Extensions;
using static RTMP.SimpleClient.Helpers.RTMPMessageHelper;

namespace RTMP.SimpleClient.Helpers
{
    /// <summary>
    /// RTMP 分块助手类，处理消息的分块和组装
    /// </summary>
    public static class RTMPChunkHelper
    {
        static readonly int MAXHEADER_SIZE = 18; // 最大头部长度（3字节基本头 + 11字节消息头 + 4字节扩展时间戳）

        #region 简化的分块接口

        /// <summary>
        /// 将完整数据包分块为RTMP数据流
        /// </summary>
        /// <param name="dataPacket">完整的数据包（包含AMF消息载荷）</param>
        /// <param name="chunkSize">分块大小</param>
        /// <param name="messageType">消息类型</param>
        /// <param name="timestamp">时间戳</param>
        /// <param name="streamId">流ID</param>
        /// <param name="chunkStreamId">分块流ID</param>
        /// <returns>分好chunk的完整数据流</returns>
        public static byte[] ChunkDataPacket(
            byte[] dataPacket,
            int chunkSize = DEFAULT_CHUNK_SIZE,
            RTMPMessageHelper.RTMPMessageType messageType = RTMPMessageHelper.RTMPMessageType.AMF0Command,
            uint timestamp = 0,
            uint streamId = 0,
            int chunkStreamId = DEFAULT_CHUNK_STREAM_ID)
        {
            if (dataPacket == null)
                throw new ArgumentNullException(nameof(dataPacket));

            // 使用GetChunkData扩展方法进行分块处理
            return dataPacket.GetChunkData(messageType, chunkStreamId, chunkSize);
        }

        /// <summary>
        /// 为AMF0命令创建分块数据流（便捷方法）
        /// </summary>
        /// <param name="commandData">AMF0命令数据</param>
        /// <param name="chunkSize">分块大小</param>
        /// <returns>分好chunk的AMF0命令数据流</returns>
        public static byte[] ChunkAMF0Command(byte[] commandData, int chunkSize = DEFAULT_CHUNK_SIZE)
        {
            return ChunkDataPacket(
                commandData,
                chunkSize,
                RTMPMessageHelper.RTMPMessageType.AMF0Command,
                timestamp: 0,
                streamId: 0,
                chunkStreamId: DEFAULT_CHUNK_STREAM_ID);
        }

        /// <summary>
        /// 为控制消息创建分块数据流（便捷方法）
        /// </summary>
        /// <param name="controlData">控制消息数据</param>
        /// <param name="messageType">控制消息类型</param>
        /// <param name="chunkSize">分块大小</param>
        /// <returns>分好chunk的控制消息数据流</returns>
        public static byte[] ChunkControlMessage(
            byte[] controlData,
            RTMPMessageHelper.RTMPMessageType messageType,
            int chunkSize = DEFAULT_CHUNK_SIZE)
        {
            return ChunkDataPacket(
                controlData,
                chunkSize,
                messageType,
                timestamp: 0,
                streamId: 0,
                chunkStreamId: 2); // 控制消息通常使用chunk stream ID 2
        }

        #endregion

        #region 常量定义

        /// <summary>
        /// 默认分块大小（128字节）
        /// </summary>
        public const int DEFAULT_CHUNK_SIZE = 128;

        /// <summary>
        /// 默认分块流ID
        /// </summary>
        public const int DEFAULT_CHUNK_STREAM_ID = 3;

        /// <summary>
        /// 分块格式枚举
        /// </summary>
        public enum ChunkFormat : byte
        {
            /// <summary>
            /// 格式0：完整头部（11字节消息头）
            /// </summary>
            Format0 = 0x00,

            /// <summary>
            /// 格式1：可变大小头部（7字节消息头）
            /// </summary>
            Format1 = 0x40,

            /// <summary>
            /// 格式2：时间戳头部（3字节消息头）
            /// </summary>
            Format2 = 0x80,

            /// <summary>
            /// 格式3：无头部（继续之前的分块）
            /// </summary>
            Format3 = 0xC0
        }

        #endregion

        #region  核心分块逻辑

        public static byte[] GetSetChunkSizePayload(int chunkSize)
        {
            Span<byte> buffer = stackalloc byte[128];
            Span<byte> span = buffer;
            int basicHeaderLength = WriteBasicHeader(buffer, ChunkFormat.Format0, 2); // chunk stream ID 2 用于控制消息
            buffer = buffer[basicHeaderLength..];
            int messageLength = WriteMessageHeader(buffer, ChunkFormat.Format0, 0, 4, (byte)RTMPMessageHelper.RTMPMessageType.SetChunkSize, 0); //完整的11字节消息头
            buffer = buffer[messageLength..];
            buffer[0..4].WriteBigEndianValue((uint)chunkSize, 4);

            return span[..(basicHeaderLength + messageLength + 4)].ToArray();
        }

        public static byte[] GetChunkData(this byte[] payload, RTMPMessageType messageType, int CSID, int chunkSize = DEFAULT_CHUNK_SIZE)
        {
            if (payload == null)
                throw new ArgumentNullException(nameof(payload));
            if (CSID < 2 || CSID > 65599)
                throw new ArgumentOutOfRangeException(nameof(CSID), "Chunk stream ID must be between 2 and 65599");
            if (chunkSize < 1 || chunkSize > 65536)
                throw new ArgumentOutOfRangeException(nameof(chunkSize), "Chunk size must be between 1 and 65536");

            //分片逻辑,按chunksize分片
            Span<byte> span = payload.AsSpan();
            bool isFirstChunk = true;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(payload.Length + (payload.Length / chunkSize + 1) * 18); // 预留足够空间
            Span<byte> bufferSpan = buffer.AsSpan();
            int offset = 0;
            int totalMessageLength = payload.Length; // 保存总消息长度

            while (span.Length > 0)
            {
                int chunkSizeToWrite = Math.Min(chunkSize, span.Length);
                int chunkPackSize = 0;
                if (isFirstChunk)
                    chunkPackSize = WriteChunkData(bufferSpan, span[..chunkSizeToWrite], ChunkFormat.Format0, messageType, CSID, totalMessageLength);
                else if (chunkSizeToWrite == chunkSize) //复用头
                {
                    chunkPackSize = WriteChunkData(bufferSpan, span[..chunkSizeToWrite], ChunkFormat.Format3, messageType, CSID, totalMessageLength);
                }
                else //增量头
                {
                    chunkPackSize = WriteChunkData(bufferSpan, span[..chunkSizeToWrite], ChunkFormat.Format1, messageType, CSID, totalMessageLength);
                }

                offset += chunkPackSize;

                bufferSpan = bufferSpan[chunkPackSize..];
                span = span[chunkSizeToWrite..];

                isFirstChunk = false;
            }

            try
            {
                bufferSpan = buffer.AsSpan();
                return bufferSpan[..offset].ToArray();
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
        }

        private static int WriteChunkData(
            Span<byte> buffer,
            Span<byte> payload,
            ChunkFormat chunkFormat,
            RTMPMessageType messageType,
            int CSID,
            int totalMessageLength)
        {
            //开始单个chunk的写入
            int headerLength = WriteBasicHeader(buffer, chunkFormat, CSID);
            int messageHeaderLength = WriteMessageHeader(buffer[headerLength..], chunkFormat, 0, (uint)totalMessageLength, (byte)messageType, 0);
            int payloadLength = WritePayload(buffer[(headerLength + messageHeaderLength)..], payload, payload.Length);

            return headerLength + messageHeaderLength + payloadLength;
        }

        private static int WritePayload(Span<byte> buffer, Span<byte> payload, int expectedLength)
        {
            if (payload.Length != expectedLength)
                throw new ArgumentException($"Payload length ({payload.Length}) must be equal to expected length ({expectedLength}).");

            payload.CopyTo(buffer);
            return payload.Length;
        }

        private static int WriteMessageHeader(Span<byte> buffer, ChunkFormat format, uint timestamp, uint messageLength, byte messageTypeId, uint messageStreamId)
        {
            switch (format)
            {
                case ChunkFormat.Format0:
                    buffer[0..3].WriteBigEndianValue(timestamp, 3);
                    buffer[3..6].WriteBigEndianValue(messageLength, 3);
                    buffer[6] = messageTypeId;
                    buffer[7..11].WriteLittleEndianValue(messageStreamId);
                    return 11; // 完整头部长度为11字节
                case ChunkFormat.Format1:
                    buffer[0..3].WriteBigEndianValue(timestamp, 3);
                    buffer[3..6].WriteBigEndianValue(messageLength, 3);
                    buffer[6] = messageTypeId;
                    return 7; // 可变大小头部长度为7字节
                case ChunkFormat.Format2:
                    // 写入时间戳头部（3字节）
                    buffer[0..3].WriteBigEndianValue(timestamp, 3);
                    return 3;
                case ChunkFormat.Format3:
                    // 无头部
                    return 0;
            }

            throw new InvalidOperationException();
        }

        private static int WriteBasicHeader(Span<byte> buffer, ChunkFormat format, int chunkStreamId)
        {
            int byteCount = GetChunkStreamIdByteCount(chunkStreamId);
            if (byteCount == 1)
            {
                buffer[0] = (byte)((byte)format | (chunkStreamId & 0x3F));
                return 1; // 基本头部长度为1字节
            }

            if (byteCount == 2)
            {
                buffer[0] = (byte)((byte)format | 0x00); // chunk stream ID 0表示后续1字节
                buffer[1] = (byte)(chunkStreamId - 64);
                return 2; // 基本头部长度为2字节
            }

            if (byteCount == 3)
            {
                buffer[0] = (byte)((byte)format | 0x01); // chunk stream ID 1表示后续2字节
                int csId = chunkStreamId - 64;
                buffer[1] = (byte)(csId & 0xFF);
                buffer[2] = (byte)((csId >> 8) & 0xFF);
                return 3; // 基本头部长度为3字节
            }

            throw new ArgumentOutOfRangeException(nameof(chunkStreamId), "Invalid chunk stream ID or buffer too small");
        }

        /// <summary>
        /// 获取分块流ID所需的字节数
        /// </summary>
        /// <param name="chunkStreamId"></param>
        /// <returns></returns>
        /// <exception cref="ArgumentOutOfRangeException"></exception>
        private static int GetChunkStreamIdByteCount(int chunkStreamId)
        {
            if (chunkStreamId < 2 || chunkStreamId > 65599)
                throw new ArgumentOutOfRangeException(nameof(chunkStreamId), "Chunk stream ID must be between 2 and 65599");

            if (chunkStreamId > 63 && chunkStreamId <= 319)
                return 2; // 2字节基本头部

            if (chunkStreamId > 319 && chunkStreamId <= 65599)
                return 3; // 3字节基本头部

            return 1; // 1字节基本头部
        }

        #endregion 核心分块逻辑

    }
}